# Copyright 2020 - 2022 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from sklearn.model_selection import KFold  ## K折交叉验证

import os
import json
import math
import numpy as np
import torch
from monai import transforms, data
import SimpleITK as sitk
from tqdm import tqdm 
from torch.utils.data import Dataset, DataLoader
from .multiprocess_loader import get_multiprocess_dataloader
from .transformer import RandomRotate, \
    GammaTransform, RemoveLabelTransform, MirrorTransform, Compose, NumpyToTensor
from .transformer.crop import RandCropByPosNegLabel
from .transformer.noisy_transformers import GaussianNoiseTransform, GaussianBlurTransform
from .transformer.color_transformers import ContrastAugmentationTransform, BrightnessMultiplicativeTransform
from .transformer.resample_transformers import SimulateLowResolutionTransform
from monai import transforms
import h5py

train_transform = [
    RandomRotate(angle_spectrum=30, execution_probability=0.2),
    RemoveLabelTransform(-1, 0),
    # RandCropByPosNegLabel([96, 96, 96], pos=1.0, neg=1.0, num_samples=2),
    GaussianNoiseTransform(noise_variance=(0, 0.1), execution_probability=0.1),
    GaussianBlurTransform(blur_sigma=(0.5, 1), execution_probability=0.1, p_per_channel=0.5, different_sigma_per_channel=True),
    BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), execution_probability=0.15),
    ContrastAugmentationTransform(contrast_range=(0.75, 1.25), execution_probability=0.15),
    SimulateLowResolutionTransform(zoom_range=(0.5, 1.0), per_channel=True, p_per_channel=0.5,
                                   order_downsample=0, order_upsample=3, execution_probability=0.1),
    GammaTransform((0.7, 1.5), True, True, execution_probability=0.1),
    GammaTransform((0.7, 1.5), False, True, execution_probability=0.1),
    # MirrorTransform((0, 1, 2)),
    NumpyToTensor(),
]

train_transform = Compose(train_transform)

val_transform = [
    RemoveLabelTransform(-1, 0),
    NumpyToTensor()
]
val_transform = Compose(val_transform)

class Dataset3d(Dataset):
    """
    """
    def __init__(
        self,
        image_paths,
        label_paths,
        has_tqdm=True
    ) -> None:
        """
        Args:
            data: input data to load and transform to generate dataset for model.
            transform: transforms to execute operations on input data.
            cache_num: number of items to be cached. Default is `sys.maxsize`.
                will take the minimum of (cache_num, data_length x cache_rate, data_length).
            cache_rate: percentage of cached data in total, default is 1.0 (cache all).
                will take the minimum of (cache_num, data_length x cache_rate, data_length).
            num_workers: the number of worker threads to use.
                If 0 a single thread will be used. Default is 0.
        """

        super(Dataset3d, self).__init__()
        self.image_paths = image_paths
        self.label_paths = label_paths

    def _load_cache_item(self, image_path, label_path):
        img = sitk.ReadImage(image_path)
        img = sitk.GetArrayFromImage(img)
        # img = img.transpose(1, 0, 2)
        label = sitk.ReadImage(label_path)
        label = sitk.GetArrayFromImage(label)
        start = 0
        end = -1
        for i in range(label.shape[0]):
            if label[i].sum() != 0:
                start = i
                break
        for j in range(label.shape[0]-1, 0, -1):
            if label[j].sum() != 0:
                end = j
                break
        img = img[start:end+1]
        label = label[start:end+1]

        # print("start is {}, end is {}".format(start, end))
        # label = label.transpose(1, 0, 2)
        return img, label

    def __getitem__(self, index):

        image, label = self._load_cache_item(self.image_paths[index], self.label_paths[index])

        image = np.expand_dims(image, axis=0)
        image = (image - np.mean(image)) / np.std(image)

        return image.astype(np.float32), label.astype(np.uint8)

    def __len__(self):
        return len(self.image_paths)


def get_loader(train_path, val_path):

    train_ds = Dataset3d(train_path["image"], train_path["label"])

    train_loader = get_multiprocess_dataloader(train_ds, train_transform=train_transform, batch_size=1, num_process=8)
    
    val_ds = Dataset3d(val_path["image"], val_path["label"])
    
    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False)

    loader = [train_loader, val_loader]

    return loader
